try:
    frozenset
except NameError:
    from sets import ImmutableSet as frozenset
from StringIO import StringIO

from genshi.core import Attrs, Markup, Namespace, QName, Stream
from genshi.core import DOCTYPE, START, END, TEXT, PI, COMMENT
from genshi.template.base import TemplateSyntaxError
from genshi.template.markup import MarkupTemplate

import html5lib
import html5lib.filters.optionaltags
from html5lib.treebuilders import getTreeBuilder
from html5lib.treewalkers import getTreeWalker

__all__ = [
'HTML5Parser', 'HTML5', 'HTML5FragmentParser', 'HTML5Fragment', \
'HTML5LibFilter', 'HTML5OptionalTagsStripper', 'HTML5Serializer',
'HTML5Template', 'HTML5FragmentTemplate', 'HTML5FragmentTemplateFactory']
__docformat__ = 'restructuredtext en'

HTML5_NAMESPACE_URI = 'http://www.w3.org/1999/xhtml'
HTML5_NAMESPACE = Namespace(HTML5_NAMESPACE_URI)

class HTML5LibAdapter(object):
    def __init__(self, source, generate_empty_events=False, filename=None):
        self.source = source
        self.generate_empty_events = generate_empty_events
        self.filename = filename

    def __iter__(self):
        pos = (self.filename, -1, -1)
        text = None
        for token in self.source:
            type = token["type"]

            if type in ('Characters', 'SpaceCharacters'):
                if text is None:
                    text = u''
                text += token["data"]
                continue # skip text yielding below

            elif type in ('ParseError', 'SerializeError'):
                continue # skip text yielding below

            elif text is not None:
                yield TEXT, text, pos
                text = None


            if type in ('StartTag', 'EmptyTag'):
                tag_name = HTML5_NAMESPACE[token["name"]]
                attrs = Attrs([(QName(name), value) for name, value in token["data"]])
                if self.generate_empty_events and type == 'EmptyTag':
                    yield EMPTY, (tag_name, attrs), pos
                else:
                    yield START, (tag_name, attrs), pos
                    if type == 'EmptyTag':
                        yield END, tag_name, pos

            elif type == 'EndTag':
                yield END, HTML5_NAMESPACE[token["name"]], pos

            elif type == 'Comment':
                yield COMMENT, token["data"], pos

            elif type == 'Doctype':
                yield DOCTYPE, (token["name"], None, None), pos

            else:
                raise Exception # TODO

        if text is not None:
            yield TEXT, text, pos
            text = None

class BaseHTML5Parser(object):
    """Base class of parsers for HTML input, based on `html5lib`.
    
    This class provides the same interface for generating stream events as
    `XMLParser`.
    
    The parsing is initiated by iterating over the parser object.
    """
    
    def __init__(self, source, innerHTML=False, container='div', filename=None, encoding=None):
        """Initialize the parser for the given HTML input.
        
        :param source: the HTML text as a file-like object
        :param innerHTML: whether this is a fragment or document parser
        :param container: name of the element we're setting the innerHTML property; only used if innerHTML is True
        :param filename: the name of the file, if known
        :param encoding: encoding of the file; ignored if the input is unicode
        """
        self.source = source
        self.innerHTML = innerHTML
        container = QName(container)
        if container in HTML5_NAMESPACE:
            container = container.localname
        self.container = container
        self.filename = filename
        self.encoding = encoding
        self.parser = html5lib.HTMLParser(tree=getTreeBuilder('simpletree'))

    def parse(self):
        """Generator that parses the HTML source, yielding markup events.
        
        :return: a markup event stream
        """
        tree = self.parser._parse(self.source, self.innerHTML, container=self.container, encoding=self.encoding)
        return HTML5LibAdapter(getTreeWalker('simpletree')(tree), generate_empty_events=False, filename=self.filename)

    def __iter__(self):
        return iter(self.parse())

class HTML5Parser(BaseHTML5Parser):
    """Parser for HTML input based on `html5lib`.
    
    This class provides the same interface for generating stream events as
    `XMLParser`.
    
    The parsing is initiated by iterating over the parser object:
    
    >>> parser = HTML5Parser(StringIO('<UL compact><LI>Foo</UL>'))
    >>> for kind, data, pos in parser:
    ...     print kind, repr(data)
    START (QName(u'html'), Attrs())
    START (QName(u'head'), Attrs())
    END QName(u'head')
    START (QName(u'body'), Attrs())
    START (QName(u'ul'), Attrs([(QName(u'compact'), '')]))
    START (QName(u'li'), Attrs())
    TEXT u'Foo'
    END QName(u'li')
    END QName(u'ul')
    END QName(u'body')
    END QName(u'html')
    """
    
    def __init__(self, source, filename=None, encoding=None):
        """Initialize the parser for the given HTML input.
        
        :param source: the HTML text as a file-like object
        :param filename: the name of the file, if known
        :param encoding: encoding of the file; ignored if the input is unicode
        """
        BaseHTML5Parser.__init__(self, source, False, filename=filename, encoding=encoding)

def HTML5(text, encoding=None):
    """Parse the given HTML source and return a markup stream.
    
    Unlike with `HTML5Parser`, the returned stream is reusable, meaning it can be
    iterated over multiple times:
    
    >>> html = HTML5('<body><h1>Foo</h1></body>')
    >>> print html
    <html xmlns="http://www.w3.org/1999/xhtml"><head/><body><h1>Foo</h1></body></html>
    >>> print html.select('body/h1')
    <h1 xmlns="http://www.w3.org/1999/xhtml">Foo</h1>
    >>> print html.select('body/h1/text()')
    Foo
    
    :param text: the HTML source
    :return: the parsed XML event stream
    """
    return Stream(list(HTML5Parser(StringIO(text), encoding=encoding)))

class HTML5FragmentParser(BaseHTML5Parser):
    """Parser for HTML input based on `html5lib`.
    
    This class provides the same interface for generating stream events as
    `XMLParser`.
    
    The parsing is initiated by iterating over the parser object:
    
    >>> parser = HTML5FragmentParser(StringIO('<UL compact><LI>Foo</UL>'), 'div')
    >>> for kind, data, pos in parser:
    ...     print kind, repr(data)
    START (QName(u'ul'), Attrs([(QName(u'compact'), '')]))
    START (QName(u'li'), Attrs())
    TEXT u'Foo'
    END QName(u'li')
    END QName(u'ul')
    """
    
    def __init__(self, source, container='div', filename=None, encoding=None):
        """Initialize the parser for the given HTML input.
        
        :param source: the HTML text as a file-like object
        :param container: name of the element we're setting the innerHTML property
        :param filename: the name of the file, if known
        :param encoding: encoding of the file; ignored if the input is unicode
        """
        HTML5Parser.__init__(self, source, True, container=container, filename=filename, encoding=encoding)

def HTML5Fragment(text, container='div', encoding=None):
    """Parse the given HTML fragment and return a markup stream.
    
    Unlike with `HTML5FragmentParser`, the returned stream is reusable, meaning
    it can be iterated over multiple times:
    
    >>> section = HTML5Fragment('<h1>Foo</h1><p>Bar', 'section')
    >>> print section
    <h1 xmlns="http://www.w3.org/1999/xhtml">Foo</h1><p xmlns="http://www.w3.org/1999/xhtml">Bar</p>
    >>> print section.select('h1')
    <h1 xmlns="http://www.w3.org/1999/xhtml">Foo</h1>
    >>> print section.select('h1/text()')
    Foo
    
    :param text: the HTML source
    :param container: name of the element we're setting the innerHTML property
    :return: the parsed XML event stream
    """
    return Stream(list(HTML5FragmentParser(StringIO(text), container=container, encoding=encoding)))

def HTML5LibFilter(filter_):
    def wrapper(source):
        return HTML5LibAdapter(filter_(getTreeWalker('genshi')(source)))
    return wrapper

HTML5OptionalTagsStripper = HTML5LibFilter(html5lib.filters.optionaltags.Filter)

class HTML5Serializer(object):
    def __init__(self, doctype=('html', None, None), serializer=None, **kwargs):
        self.preamble = []
        if doctype:
            self.preamble.append((DOCTYPE, doctype, (None, -1, -1)))
        self.filters = []
        if serializer is not None:
            self.serializer = serializer
        else:
            # Disable optional tags omitting by default
            kwargs.setdefault('omit_optional_tags', False)
            self.serializer = HTMLSerializer(**kwargs)

    def _ensure_no_duplicate_doctype(stream):
        have_doctype = False
        for kind, data, pos in stream:
            if kind is not DOCTYPE or not have_doctype:
                yield kind, data, pos
                have_doctype = True
    _ensure_no_duplicate_doctype = staticmethod(_ensure_no_duplicate_doctype)

    def __call__(self, stream):
        stream = chain(self.preamble, stream)
        for filter_ in self.filters:
            stream = filter_(stream)

        stream = self._ensure_no_duplicate_doctype(stream)

        for markup in self.serializer.serialize(getTreeWalker('genshi')(stream)):
            yield Markup(markup)

class BaseHTML5Template(MarkupTemplate):
    """Implementation of the template language for HTML-based templates.
    
    This is an abstract class. Concrete implementations have to implement
    the ``_create_parser`` method.
    """

    def __init__(self, source, basedir=None, filename=None, loader=None,
                  encoding=None):
        MarkupTemplate.__init__(self, source, basedir=basedir, filename=filename,
                                loader=loader, encoding=encoding)

    def _create_parser(self, source, encoding):
        raise NotImplementedError
    
    def _parse(self, source, encoding):
        if not isinstance(source, Stream):
            source = self._create_parser(source, encoding)
            source = Stream(source)
        source = source | self._to_xml
        return MarkupTemplate._parse(self, source, encoding)

    def _to_xml(self, stream):
        """Transforms elements or attributes from HTML namespace (or no namespace) \
        whose name starts with py: or pi_ into XML Template directives.
        
        Also transforms elements from HTML namespace (or no namespace) whose name
        starts with xi: or xi_ into XInclude elements.
        
        Templates can be included using <link rel="genshi.include" href="...">,
        <script type="text/html" src="..."></script> or
        <object classid="genshi.include" src="...">fallback</object>.
        
        Python code can be include within <script type="text/x-python">...</script>
        Recognized type values are text/x-python, text/python, application/x-python
        and application/python.
        """
        in_include = 0
        in_script = 0
        
        for kind, data, pos in stream:
            if kind is START:
                tag, attrs = data
                if not tag.namespace or tag in HTML5_NAMESPACE:
                    if tag.localname.startswith('py:') or tag.localname.startswith('py_'):
                        tag = self.DIRECTIVE_NAMESPACE[tag.localname[3:]]
                    elif tag.localname.startswith('xi:') or tag.localname.startswith('xi_'):
                        tag = self.XINCLUDE_NAMESPACE[tag.localname[3:]]
                    elif tag.localname == 'link' and attrs.get('rel') == 'genshi.include':
                        in_include += 1
                        tag = self.XINCLUDE_NAMESPACE['include']
                        attrs = attrs - 'rel'
                    elif tag.localname == 'script':
                        if attrs.get('type') == 'text/html' and 'src' in attrs:
                            in_include += 1
                            tag = self.XINCLUDE_NAMESPACE['include']
                            attrs = attrs - ('type', 'src') | [('href', attrs['src'])]
                        elif attrs.get('type') in ('text/x-python', 'text/python', 'application/x-python', 'application/python'):
                            in_script += 1
                            continue # Script will be emitted at "end tag"
                    elif tag.localname == 'object' and attrs.get('classid') == 'genshi.include':
                        in_include += 1
                        yield (START, (self.XINCLUDE_NAMESPACE['include'], Attrs([('href', attrs.get('src'))])), pos)
                        tag = self.XINCLUDE_NAMESPACE['fallback']
                        attrs = Attrs()
                new_attrs = []
                for name, value in attrs:
                    if not name.namespace or name in HTML5_NAMESPACE:
                        if name.localname.startswith('py:') or name.localname.startswith('py_'):
                            name = self.DIRECTIVE_NAMESPACE[name.localname[3:]]
                    new_attrs.append((name, value))
                yield kind, (tag, Attrs(new_attrs)), pos

            elif kind is END:
                if not data.namespace or data in HTML5_NAMESPACE:
                    if data.localname.startswith('py:') or data.localname.startswith('py_'):
                        data = self.DIRECTIVE_NAMESPACE[data.localname[3:]]
                    elif data.localname.startswith('xi:') or data.localname.startswith('xi_'):
                        data = self.XINCLUDE_NAMESPACE[data.localname[3:]]
                    elif in_include and data.localname in ('link', 'script', 'object'):
                        if data.localname == 'object':
                            yield (END, self.XINCLUDE_NAMESPACE['fallback'], pos)
                        data = self.XINCLUDE_NAMESPACE['include']
                        in_include -= 1
                    elif in_script and data.localname == 'script':
                        in_script -= 1
                yield kind, data, pos

            elif in_script:
                if kind is TEXT:
                    yield (PI, ('python', data), pos)
                else:
                    raise TemplateSyntaxError('Inline Python code must contain '
                                              'only text (found %s)' % kind,
                                              self.filepath, *pos[1:])

            elif in_include:
                raise TemplateSyntaxError('Include must be empty '
                                          '(found %s)' % kind,
                                          self.filepath, *pos[1:])

            else:
                yield kind, data, pos

    def _include(self, stream, ctxt):
        """Internal stream filter that performs inclusion of external
        template files.
        """
        openElements = [False]
        for event in stream:
            if event[0] is INCLUDE:
                href, fallback = event[1]
                if not isinstance(href, basestring):
                    parts = []
                    for subkind, subdata, subpos in self._eval(href, ctxt):
                        if subkind is TEXT:
                            parts.append(subdata)
                    href = u''.join([x for x in parts if x is not None])
                try:
                    tmpl = self.loader.load(href, relative_to=event[2][0],
                                            cls=HTML5FragmentTemplateFactory(openElements[-1]))
                    for event in tmpl.generate(ctxt):
                        yield event
                except TemplateNotFound:
                    if fallback is None:
                        raise
                    for filter_ in self.filters:
                        fallback = filter_(iter(fallback), ctxt)
                    for event in fallback:
                        yield event
            else:
                if event[0] is START:
                    openElements.append(event[1][0])
                elif event[0] is END:
                    openElements.pop()
                yield event

class HTML5Template(BaseHTML5Template):
    """Implementation of the template language for HTML-based templates.
    
    >>> tmpl = HTML5Template('''<ul>
    ...   <li py:for="item in items">${item}</li>
    ... </ul>''')
    >>> print tmpl.generate(items=[1, 2, 3])
    <html xmlns="http://www.w3.org/1999/xhtml"><head/><body><ul>
      <li>1</li><li>2</li><li>3</li>
    </ul></body></html>
    """

    def __init__(self, source, basedir=None, filename=None, loader=None,
                  encoding=None):
        BaseHTML5Template.__init__(self, source, basedir=basedir, filename=filename,
                                  loader=loader, encoding=encoding)

    def _create_parser(self, source, encoding):
        source = HTML5Parser(source, filename=self.filename, encoding=encoding)

class HTML5FragmentTemplate(BaseHTML5Template):
    """Implementation of the template language for HTML-based templates.
    
    >>> tmpl = HTML5Template('''<ul>
    ...   <li py:for="item in items">${item}</li>
    ... </ul>''')
    >>> print tmpl.generate(items=[1, 2, 3])
    <html xmlns="http://www.w3.org/1999/xhtml"><head/><body><ul>
      <li>1</li><li>2</li><li>3</li>
    </ul></body></html>
    """

    def __init__(self, source, basedir=None, filename=None, loader=None,
                  encoding=None, container='div'):
        BaseHTML5Template.__init__(self, source, basedir=basedir, filename=filename,
                                  loader=loader, encoding=encoding)
        self.container = container

    def _create_parser(self, source, encoding):
        source = HTML5FragmentParser(source, container=self.container, filename=self.filename, encoding=encoding)

class HTML5FragmentTemplateFactory(object):
    def __init__(self, container='div'):
        self.container = container
    
    def __call__(self, source, basedir=None, filename=None, loader=None, encoding=None):
        return HTML5FragmentTemplate(source, basedir=basedir, filename=filename,
                                     loader=loader, encoding=encoding, container=self.container)
